#############################################################################

#     IPW method for mediation analysis               
#     Applicable situation:
#         1.The outcome is dichotomous data;
#         2.An exposure variable and two covariates: dichotomous data;
#         2.Two parallel mediators.

#############################################################################
IPW<-function(Data,n){
  Data$id<-c(1:n)
  
  #Fit a suitable model for the probability (density) of the mediators
  if(all(unique(Data$M1)%in%c(0,1))){
    fitM1 <- glm(M1 ~ A+C1+C2,family = binomial("logit"), data = Data)
  }else{
    fitM1 <- glm(M1 ~ A+C1+C2,family = gaussian("identity"), data = Data)
  }
  beta1<-as.character(coef(summary(fitM1))[2,1])
  beta1<-as.numeric(beta1)
  
  if(all(unique(Data$M2)%in%c(0,1))){
    fitM2 <- glm(M2 ~ A+C1+C2,family = binomial("logit"), data = Data)
  }else{
    fitM2 <- glm(M2 ~ A+C1+C2,family = gaussian("identity"), data = Data)
  }
  beta2<-as.character(coef(summary(fitM2))[2,1])
  beta2<-as.numeric(beta2)
  
  #Fit a suitable model for the outcome mean
  fitY <- glm(Y ~ A + M1 + M2 + C1 + C2,family = binomial("logit"), data = Data)
  
  #Construct the extended data set
  extdat <- data.frame(replicate = rep(1:4, times = nrow(Data)),Data[rep(Data$id, each = 4), ],a0 = NA, a1 = NA, a2 = NA)
  
  extdat1<-within(extdat,{
    a0 <- ifelse(replicate %in% c(2,4), 1-A, A)
    a1 <- ifelse(replicate %in% c(3,4), 1-A, A)
    a2<-A
  })
  
  extdat2<-within(extdat,{
    a0 <- ifelse(replicate %in% c(2,4), 1-A, A)
    a1<-A
    a2 <- ifelse(replicate %in% c(3,4), 1-A, A)
  })
  
  #Calculate regression weights
  #calculate W1 for mediator M1
  if(all(unique(Data$M1)%in%c(0,1))){
    num1 <- with(extdat1,
                 dbinom(M1, size = 1, prob = predict(fitM1, newdata = within(extdat1, A <- a1),type = "response")))
    
    denom1 <- with(extdat1,
                   dbinom(M1, size = 1, prob = predict(fitM1, newdata = within(extdat1, A <- a2),type = "response")))
    
    extdat1$W1 <- num1/denom1
  }else{
    num1 <- with(extdat1,
                 dnorm(M1, mean = predict(fitM1, newdata = within(extdat1, A <- a1),type = "response"),
                       sd = sqrt(summary(fitM1)$dispersion)))
    
    denom1 <- with(extdat1,
                   dnorm(M1, mean = predict(fitM1, newdata = within(extdat1, A <- a2),type = "response"),
                         sd = sqrt(summary(fitM1)$dispersion)))
    
    extdat1$W1 <- num1/denom1
  }
  
  #calculate W2 for mediator M2
  if(all(unique(Data$M2)%in%c(0,1))){
    num2 <- with(extdat2,
                 dbinom(M2, size = 1, prob = predict(fitM2, newdata = within(extdat2, A <- a2),type = "response")))
    
    denom2 <- with(extdat2,
                   dbinom(M2, size = 1, prob = predict(fitM2, newdata = within(extdat2, A <- a1),type = "response")))
    
    extdat2$W2 <- num2/denom2
  }else{
    num2 <- with(extdat2,
                 dnorm(M2, mean = predict(fitM2, newdata = within(extdat2, A <- a2),type = "response"),
                       sd = sqrt(summary(fitM2)$dispersion)))
    
    denom2 <- with(extdat2,
                   dnorm(M2, mean = predict(fitM2, newdata = within(extdat2, A <- a1),type = "response"),
                         sd = sqrt(summary(fitM2)$dispersion)))
    
    extdat2$W2 <- num2/denom2
  }
  
  #Obtain population-average component effects(rather than effects conditional on C)
  fitA <- glm(A ~ C1+C2,family = binomial("logit"), data = Data)
  
  #Calculate the final weight
  finalextdat<-rbind(extdat1,extdat2$W2)
  names(finalextdat)[13]<-c("W2")
  finalextdat<-within(finalextdat,{
    W<-W1 * W2/dbinom(A, size = 1, prob = predict(fitA, newdata = finalextdat,type = "response"))
  })
  
  #Fit a natural effect model
  fitNEM <- glm(Y ~ A + M1 + M2 + C1 + C2,
                family = binomial("logit"), data = finalextdat, weights = W)
  
  #Direct effect
  DE<-as.character(coef(summary(fitNEM))[2,1])
  DE<-as.numeric(DE)
  
  #Indirect effect
  theta1<-as.character(coef(summary(fitNEM))[3,1])
  theta1<-as.numeric(theta1)
  theta2<-as.character(coef(summary(fitNEM))[4,1])
  theta2<-as.numeric(theta2)
  IE1<-theta1*beta1
  IE2<-theta2*beta2
  
  result<-data.frame(DE,IE1,IE2)
  return(result)
    
}